#include "timer.h"

#include <linux/proc_fs.h>
#include <linux/sched.h>
#include <linux/timer.h>
#include <linux/completion.h>
#include <linux/uaccess.h> // Required for copy_to_user

#include "log.h"

#define BYTES_PER_LINE 128 // Increased buffer size for safety

int timer_interval = 1000;

struct jit_timer_context {
    struct timer_list timer;
    char *buf;
    char *p_buf; // current pointer in buf
    size_t buf_size;
    int loops;
    struct completion done;
    ulong jiffies;
};

static void jit_timer_work(struct timer_list *timer) {
    struct jit_timer_context *ctx = container_of(timer, struct jit_timer_context, timer);
    
    if (ctx) {
        ulong now = jiffies;
        size_t remaining_buf_size = ctx->buf_size - (ctx->p_buf - ctx->buf);

        // Use snprintf to safely write to the buffer
        ctx->p_buf += snprintf(ctx->p_buf, remaining_buf_size, "%lu\t%lu\t%d\t%d\t%d\t%s\n",
            now, now - ctx->jiffies, in_interrupt() ? 1 : 0, task_pid_nr(current), smp_processor_id(), current->comm);

        if (--ctx->loops > 0) {
            ctx->jiffies = now;
            ctx->timer.expires += msecs_to_jiffies(timer_interval);
            mod_timer(&ctx->timer, ctx->timer.expires);
        } else {
            // Last loop, signal completion
            complete(&ctx->done);
        }
    } else {
        LOGW("timer data is invalid");
    }
}

// This is our new custom read function
static ssize_t jit_timer_read(struct file *filp, char __user *ubuf, size_t count, loff_t *f_pos) {
    struct jit_timer_context *ctx;
    ssize_t ret = 0;
    size_t data_len;

    LOGI("jit_timer_read enter, f_pos=%lld", *f_pos);

    // We only allow reading from the beginning. Subsequent reads will return 0 (EOF).
    if (*f_pos > 0) {
        return 0;
    }

    // 1. Allocate context and buffer
    ctx = kzalloc(sizeof(*ctx), GFP_KERNEL);
    if (!ctx) {
        LOGE("out of memory for context");
        return -ENOMEM;
    }
    
    ctx->loops = 5;
    ctx->buf_size = (BYTES_PER_LINE * (ctx->loops + 1)); // +1 for header
    ctx->buf = kzalloc(ctx->buf_size, GFP_KERNEL);
    if (!ctx->buf) {
        LOGE("out of memory for buffer");
        kfree(ctx);
        return -ENOMEM;
    }
    ctx->p_buf = ctx->buf; // Initialize current pointer
    init_completion(&ctx->done);

    // 2. Prepare header and start timer
    ulong now = jiffies;
    ctx->jiffies = now;
    ctx->p_buf += snprintf(ctx->p_buf, ctx->buf_size, "time\t\tdelta\tinirq\tpid\tcpu\tcommand\n");
    ctx->p_buf += snprintf(ctx->p_buf, ctx->buf_size - (ctx->p_buf - ctx->buf), "%lu\t%lu\t%d\t%d\t%d\t%s\n",
        now, 0L, in_interrupt() ? 1 : 0, task_pid_nr(current), smp_processor_id(), current->comm);

    timer_setup(&ctx->timer, jit_timer_work, 0);
    ctx->timer.expires = now + msecs_to_jiffies(timer_interval);
    add_timer(&ctx->timer);

    // 3. Wait for the timer to complete. This is SAFE here.
    // The process will sleep, but it's in a regular read context, not atomic.
    LOGI("waiting for completion...");
    wait_for_completion(&ctx->done);
    LOGI("done");

    // 4. Copy the collected data to the user
    data_len = ctx->p_buf - ctx->buf;
    if (copy_to_user(ubuf, ctx->buf, data_len)) {
        ret = -EFAULT;
    } else {
        ret = data_len;
        *f_pos += data_len; // Advance file position
    }

    // 5. Clean up
    kfree(ctx->buf);
    kfree(ctx);

    LOGI("jit_timer_read exit, returning %zd", ret);
    return ret;
}


// We no longer need seq_file, so we define our own proc_ops
const struct proc_ops timer_ops = {
    // We don't need a special open, the kernel default is fine.
    .proc_read = jit_timer_read,
    // We need lseek and release for the file operations to be complete.
    .proc_lseek = no_llseek,
    .proc_release = NULL, // Default is fine
};

int jit_timer_init(void) {
    proc_create("jittimer", 0644, NULL, &timer_ops);
    return 0;
}

void jit_timer_exit(void) {
    remove_proc_entry("jittimer", NULL);
}